// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <nebula/types/schema.h>
#include <nebula/types/type.h>
#include <nebula/types/util.h>
#include <nebula/util/vector.h>

namespace nebula {

    class Schema::Impl {
    public:
        Impl(FieldVector fields, Endianness endianness,
             std::shared_ptr<const KeyValueMetadata> metadata)
                : fields_(std::move(fields)),
                  endianness_(endianness),
                  name_to_index_(internal::create_name_to_index_map(fields_)),
                  metadata_(std::move(metadata)) {}

        FieldVector fields_;
        Endianness endianness_;
        std::unordered_multimap<std::string, int> name_to_index_;
        std::shared_ptr<const KeyValueMetadata> metadata_;
    };



    Schema::Schema(FieldVector fields, Endianness endianness,
                   std::shared_ptr<const KeyValueMetadata> metadata)
            : detail::FingerPrintAble(),
              impl_(new Impl(std::move(fields), endianness, std::move(metadata))) {}

    Schema::Schema(FieldVector fields, std::shared_ptr<const KeyValueMetadata> metadata)
            : detail::FingerPrintAble(),
              impl_(new Impl(std::move(fields), Endianness::Native, std::move(metadata))) {}

    Schema::Schema(const Schema &schema)
            : detail::FingerPrintAble(), impl_(new Impl(*schema.impl_)) {}

    Schema::~Schema() = default;

    std::shared_ptr<Schema> Schema::with_endianness(Endianness endianness) const {
        return std::make_shared<Schema>(impl_->fields_, endianness, impl_->metadata_);
    }

    Endianness Schema::endianness() const { return impl_->endianness_; }

    bool Schema::is_native_endian() const { return impl_->endianness_ == Endianness::Native; }

    int Schema::num_fields() const { return static_cast<int>(impl_->fields_.size()); }

    const std::shared_ptr<Field> &Schema::field(int i) const {
        DKCHECK_GE(i, 0);
        DKCHECK_LT(i, num_fields());
        return impl_->fields_[i];
    }

    const FieldVector &Schema::fields() const { return impl_->fields_; }

    bool Schema::equals(const Schema &other, bool check_metadata) const {
        if (this == &other) {
            return true;
        }

        // checks endianness equality
        if (endianness() != other.endianness()) {
            return false;
        }

        // checks field equality
        if (num_fields() != other.num_fields()) {
            return false;
        }

        if (check_metadata) {
            const auto &metadata_fp = metadata_fingerprint();
            const auto &other_metadata_fp = other.metadata_fingerprint();
            if (metadata_fp != other_metadata_fp) {
                return false;
            }
        }

        // Fast path using fingerprints, if possible
        const auto &fp = fingerprint();
        const auto &other_fp = other.fingerprint();
        if (!fp.empty() && !other_fp.empty()) {
            return fp == other_fp;
        }

        // Fall back on field-by-field comparison
        for (int i = 0; i < num_fields(); ++i) {
            if (!field(i)->equals(*other.field(i).get(), check_metadata)) {
                return false;
            }
        }

        return true;
    }

    bool Schema::equals(const std::shared_ptr<Schema> &other, bool check_metadata) const {
        if (other == nullptr) {
            return false;
        }

        return equals(*other, check_metadata);
    }

    std::shared_ptr<Field> Schema::get_field_by_name(const std::string &name) const {
        int i = get_field_index(name);
        return i == -1 ? nullptr : impl_->fields_[i];
    }

    int Schema::get_field_index(const std::string &name) const {
        return internal::lookup_name_index(impl_->name_to_index_, name);
    }

    std::vector<int> Schema::get_all_field_indices(const std::string &name) const {
        std::vector<int> result;
        auto p = impl_->name_to_index_.equal_range(name);
        for (auto it = p.first; it != p.second; ++it) {
            result.push_back(it->second);
        }
        if (result.size() > 1) {
            std::sort(result.begin(), result.end());
        }
        return result;
    }

    turbo::Status Schema::can_reference_field_by_name(const std::string &name) const {
        if (get_field_by_name(name) == nullptr) {
            return turbo::invalid_argument_error("Field named '", name,
                                                 "' not found or not unique in the schema.");
        }
        return turbo::OkStatus();
    }

    turbo::Status Schema::can_reference_fields_by_names(const std::vector<std::string> &names) const {
        for (const auto &name: names) {
            TURBO_RETURN_NOT_OK(can_reference_field_by_name(name));
        }
        return turbo::OkStatus();
    }

    FieldVector Schema::get_all_fields_by_name(const std::string &name) const {
        FieldVector result;
        auto p = impl_->name_to_index_.equal_range(name);
        for (auto it = p.first; it != p.second; ++it) {
            result.push_back(impl_->fields_[it->second]);
        }
        return result;
    }

    turbo::Result<std::shared_ptr<Schema>> Schema::add_field(
            int i, const std::shared_ptr<Field> &field) const {
        if (i < 0 || i > this->num_fields()) {
            return turbo::invalid_argument_error("Invalid column index to add field.");
        }

        return std::make_shared<Schema>(internal::AddVectorElement(impl_->fields_, i, field),
                                        impl_->metadata_);
    }

    turbo::Result<std::shared_ptr<Schema>> Schema::set_field(
            int i, const std::shared_ptr<Field> &field) const {
        if (i < 0 || i > this->num_fields()) {
            return turbo::invalid_argument_error("Invalid column index to set field.");
        }

        return std::make_shared<Schema>(
                internal::ReplaceVectorElement(impl_->fields_, i, field), impl_->metadata_);
    }

    turbo::Result<std::shared_ptr<Schema>> Schema::remove_field(int i) const {
        if (i < 0 || i >= this->num_fields()) {
            return turbo::invalid_argument_error("Invalid column index to remove field.");
        }

        return std::make_shared<Schema>(internal::DeleteVectorElement(impl_->fields_, i),
                                        impl_->metadata_);
    }

    bool Schema::has_metadata() const {
        return (impl_->metadata_ != nullptr) && (impl_->metadata_->size() > 0);
    }

    bool Schema::has_distinct_field_names() const {
        auto fields = field_names();
        std::unordered_set<std::string> names{fields.cbegin(), fields.cend()};
        return names.size() == fields.size();
    }

    turbo::Result<std::shared_ptr<Schema>> Schema::with_names(
            const std::vector<std::string> &names) const {
        if (names.size() != impl_->fields_.size()) {
            return turbo::invalid_argument_error("attempted to rename schema with ", impl_->fields_.size(),
                                                 " fields but only ", names.size(), " new names were given");
        }
        FieldVector new_fields;
        new_fields.reserve(names.size());
        auto names_itr = names.begin();
        for (const auto &field: impl_->fields_) {
            new_fields.push_back(field->with_name(*names_itr++));
        }
        return schema(std::move(new_fields));
    }

    std::shared_ptr<Schema> Schema::with_metadata(
            const std::shared_ptr<const KeyValueMetadata> &metadata) const {
        return std::make_shared<Schema>(impl_->fields_, metadata);
    }

    const std::shared_ptr<const KeyValueMetadata> &Schema::metadata() const {
        return impl_->metadata_;
    }

    std::shared_ptr<Schema> Schema::remove_metadata() const {
        return std::make_shared<Schema>(impl_->fields_);
    }

    std::string Schema::to_string(bool show_metadata) const {
        std::stringstream buffer;

        int i = 0;
        for (const auto &field: impl_->fields_) {
            if (i > 0) {
                buffer << std::endl;
            }
            buffer << field->to_string(show_metadata);
            ++i;
        }

        if (impl_->endianness_ != Endianness::Native) {
            buffer << "\n-- endianness: " << endianness_to_string(impl_->endianness_) << " --";
        }

        if (show_metadata && has_metadata()) {
            buffer << impl_->metadata_->to_string();
        }

        return buffer.str();
    }

    std::vector<std::string> Schema::field_names() const {
        std::vector<std::string> names;
        for (const auto &field: impl_->fields_) {
            names.push_back(field->name());
        }
        return names;
    }



    std::string Schema::compute_fingerprint() const {
        std::stringstream ss;
        ss << "S{";
        for (const auto &field: fields()) {
            const auto &field_fingerprint = field->fingerprint();
            if (field_fingerprint.empty()) {
                return "";
            }
            ss << field_fingerprint << ";";
        }
        ss << (endianness() == Endianness::Little ? "L" : "B");
        ss << "}";
        return ss.str();
    }

    std::string Schema::compute_metadata_fingerprint() const {
        std::stringstream ss;
        if (has_metadata()) {
            internal::append_metadata_finger_print(*metadata(), &ss);
        }
        ss << "S{";
        for (const auto &field: fields()) {
            const auto &field_fingerprint = field->metadata_fingerprint();
            ss << field_fingerprint << ";";
        }
        ss << "}";
        return ss.str();
    }

    void print_to(const Schema &s, std::ostream *os) { *os << s; }





    class SchemaBuilder::Impl {
    public:
        friend class SchemaBuilder;

        Impl(ConflictPolicy policy, Field::MergeOptions field_merge_options)
                : policy_(policy), field_merge_options_(field_merge_options) {}

        Impl(FieldVector fields, std::shared_ptr<const KeyValueMetadata> metadata,
             ConflictPolicy conflict_policy, Field::MergeOptions field_merge_options)
                : fields_(std::move(fields)),
                  name_to_index_(internal::create_name_to_index_map(fields_)),
                  metadata_(std::move(metadata)),
                  policy_(conflict_policy),
                  field_merge_options_(field_merge_options) {}

        turbo::Status add_field(const std::shared_ptr<Field> &field) {
                    DKCHECK_NE(field, nullptr);

            // Short-circuit, no lookup needed.
            if (policy_ == CONFLICT_APPEND) {
                return append_field(field);
            }

            auto name = field->name();
            constexpr int kNotFound = -1;
            constexpr int kDuplicateFound = -2;
            auto i = internal::lookup_name_index<kNotFound, kDuplicateFound>(name_to_index_, name);

            if (i == kNotFound) {
                return append_field(field);
            }

            // From this point, there's one or more field in the builder that exists with
            // the same name.

            if (policy_ == CONFLICT_IGNORE) {
                // The ignore policy is more generous when there's duplicate in the builder.
                return turbo::OkStatus();
            } else if (policy_ == CONFLICT_ERROR) {
                return turbo::invalid_argument_error("Duplicate found, policy dictate to treat as an error");
            }

            if (i == kDuplicateFound) {
                // Cannot merge/replace when there's more than one field in the builder
                // because we can't decide which to merge/replace.
                return turbo::invalid_argument_error("Cannot merge field ", name,
                                                     " more than one field with same name exists");
            }

                    DKCHECK_GE(i, 0);

            if (policy_ == CONFLICT_REPLACE) {
                fields_[i] = field;
            } else if (policy_ == CONFLICT_MERGE) {
                TURBO_MOVE_OR_RAISE(fields_[i],
                                    fields_[i]->merge_with(field, field_merge_options_));
            }

            return turbo::OkStatus();
        }

        turbo::Status append_field(const std::shared_ptr<Field> &field) {
            name_to_index_.emplace(field->name(), static_cast<int>(fields_.size()));
            fields_.push_back(field);
            return turbo::OkStatus();
        }

        void reset() {
            fields_.clear();
            name_to_index_.clear();
            metadata_.reset();
        }

    private:
        FieldVector fields_;
        std::unordered_multimap<std::string, int> name_to_index_;
        std::shared_ptr<const KeyValueMetadata> metadata_;
        ConflictPolicy policy_;
        Field::MergeOptions field_merge_options_;
    };



    SchemaBuilder::SchemaBuilder(ConflictPolicy policy,
                                 Field::MergeOptions field_merge_options) {
        impl_ = std::make_unique<Impl>(policy, field_merge_options);
    }

    SchemaBuilder::SchemaBuilder(FieldVector fields, ConflictPolicy policy,
                                 Field::MergeOptions field_merge_options) {
        impl_ = std::make_unique<Impl>(std::move(fields), nullptr, policy, field_merge_options);
    }

    SchemaBuilder::SchemaBuilder(const std::shared_ptr<Schema> &schema, ConflictPolicy policy,
                                 Field::MergeOptions field_merge_options) {
        std::shared_ptr<const KeyValueMetadata> metadata;
        if (schema->has_metadata()) {
            metadata = schema->metadata()->copy();
        }

        impl_ = std::make_unique<Impl>(schema->fields(), std::move(metadata), policy,
                                       field_merge_options);
    }

    SchemaBuilder::~SchemaBuilder() {}

    SchemaBuilder::ConflictPolicy SchemaBuilder::policy() const { return impl_->policy_; }

    void SchemaBuilder::set_policy(SchemaBuilder::ConflictPolicy resolution) {
        impl_->policy_ = resolution;
    }

    turbo::Status SchemaBuilder::add_field(const std::shared_ptr<Field> &field) {
        return impl_->add_field(field);
    }

    turbo::Status SchemaBuilder::add_fields(const FieldVector &fields) {
        for (const auto &field: fields) {
            TURBO_RETURN_NOT_OK(add_field(field));
        }

        return turbo::OkStatus();
    }

    turbo::Status SchemaBuilder::add_schema(const std::shared_ptr<Schema> &schema) {
                DKCHECK_NE(schema, nullptr);
        return add_fields(schema->fields());
    }

    turbo::Status SchemaBuilder::add_schemas(const std::vector<std::shared_ptr<Schema>> &schemas) {
        for (const auto &schema: schemas) {
            TURBO_RETURN_NOT_OK(add_schema(schema));
        }

        return turbo::OkStatus();
    }

    turbo::Status SchemaBuilder::add_metadata(const KeyValueMetadata &metadata) {
        impl_->metadata_ = metadata.copy();
        return turbo::OkStatus();
    }

    turbo::Result<std::shared_ptr<Schema>> SchemaBuilder::finish() const {
        return schema(impl_->fields_, impl_->metadata_);
    }

    void SchemaBuilder::reset() { impl_->reset(); }

    turbo::Result<std::shared_ptr<Schema>> SchemaBuilder::Merge(
            const std::vector<std::shared_ptr<Schema>> &schemas, ConflictPolicy policy) {
        SchemaBuilder builder{policy};
        TURBO_RETURN_NOT_OK(builder.add_schemas(schemas));
        return builder.finish();
    }

    turbo::Status SchemaBuilder::are_compatible(const std::vector<std::shared_ptr<Schema>> &schemas,
                                               ConflictPolicy policy) {
        return Merge(schemas, policy).status();
    }

}  // namespace nebula
